import pandas as pd
import util.lib as lib

base_column = ['date','stock','trade_code','wg','tag','wg_industry','inc_num','gt60m',
                                   'month_pct','next_month_pct' ,'close','lp','rel_lp','category']

def util_new_empty_df():
    return pd.DataFrame({},columns=base_column)

def util_built_item(k, item):
    return {'date':item.name,'trade_code':item.trade_code,
                             'stock': k,
                             'wg_industry': item.wg_industry,
                            'tag':item.tag,
                            'wg':item.wg,
                            'close':item.close,
                             'month_pct': item.month_pct,
                            'next_month_pct': item.next_month_pct,
                             'lp': item.level_point,
                             'rel_lp': item.rel_lp,
                             'inc_num': item.inc_num,
                             'gt60m': item.gt1200d,
                            'category':item.category
                            }
def util_built(k, df):
    return util_built_item(k, df.iloc[-1])


def util_get_df(v, end):
    if end=='':
        return v
    if v.iloc[0].name < end:
        return v.loc[:end]
    
    return []

# 计算相差月数
def util_cal_duraton(start, end):
    s_year,s_month = start.split("--")
    e_year,e_month = end.split("--")
    
    return (int(e_year) - int(s_year))*12 + int(e_month) - int(s_month)

# 分位点新高
def up100lp(data, end = ''):
    tmp = pd.DataFrame({},columns=base_column)

    for k,v in data.items():
        df = util_get_df(v, end)
        if len(df) < 2:
            continue

        if (df.iloc[-1].level_point >=99) and (df.iloc[-2].level_point < 99):
            new_item = util_built(k, df)
            tmp.loc[len(tmp)] = new_item
    return tmp

# 本月lp突破10的
def up10lp(data, end = ''):

    tmp = pd.DataFrame({},columns=base_column+['pre_date','duration','month_count'])
    for k,v in data.items():
        
        df = util_get_df(v, end)
        if len(df) < 2:
            continue
            
        if (df.iloc[-1].level_point >=10) and (df.iloc[-2].level_point < 10):
            new_item = util_built(k, df)
            
            # 计算之前持续几个月低于 10
            n = -1
            while len(df) > (-1 *n):
                first_item = df.iloc[n]
                if df.iloc[n-1].level_point > 10:
                    break
                n = n-1
            
            new_item['pre_date'] = first_item.name
            new_item['duration'] = util_cal_duraton(new_item['pre_date'],new_item['date'])
            new_item['month_count'] = len(df)
            
            tmp.loc[len(tmp)] = new_item
    return tmp


# 本月lp突破80的
def up80lp(data, end = ''):

    tmp = pd.DataFrame({},columns=base_column+['month_count'])
    for k,v in data.items():
        
        df = util_get_df(v, end)
        if len(df) == 0:
            continue
            
        if (df.iloc[-1].level_point >=80) and (df.iloc[-2].level_point < 80):
            new_item = util_built(k, df)
            new_item['month_count'] = len(df)
            tmp.loc[len(tmp)] = new_item
    return tmp


# 本月lp跌破10的
def down10lp(data, end = ''):
    tmp = util_new_empty_df()
    for k,v in data.items():

        df = util_get_df(v, end)
        if len(df) == 0:
            continue
            
        if (df.iloc[-1].level_point < 10) and (df.iloc[-2].level_point >= 10):
            tmp.loc[len(tmp)] = util_built(k, df)
    return tmp

# 月涨跌统计
def month_statistic(data, end=''):
    ret = pd.DataFrame({},columns=base_column+['pre_lp','de_max','pct_sum'])
    for k,v in data.items():

        df = util_get_df(v, end)
        if len(df) < 2:
            continue
        
        item = util_built(k, df)
        item['de_max'] = df['month_pct'].min()

        item['pre_lp'] =df.iloc[-2].level_point

        # 计算连续上涨或下跌周期
        if (item['month_pct'] >= 0):
            # 计算上涨
            count = 1
            in_sum = item['month_pct']
            try:
                while(df.iloc[-1*count - 1].month_pct >= 0):
                    in_sum += df.iloc[-1*count - 1].month_pct
                    count += 1
            except IndexError:
                pass

            item['pct_sum'] = in_sum
        else:
            # 计算下跌
            count = 1
            de_sum = item['month_pct']
            try:
                while(df.iloc[-1*count - 1].month_pct < 0):
                    de_sum += df.iloc[-1*count - 1].month_pct
                    count += 1
            except IndexError:
                pass

            item['pct_sum'] = de_sum

        ret.loc[len(ret)] = item

    ret.drop(columns=["wg","trade_code"], inplace=True)

    return ret

def lp_decrease(data):
    ret = pd.DataFrame({},columns=base_column+['lp_decrease','lp_duration','duration'])
    for k,df in data.items():
        lpClose = 0
        lpDuration = 0

        for index, row in df.iterrows():
            if row.level_point > 10:
                lpClose = 0
                lpDuration = 0
                continue
            else: 
                lpDuration += 1

            # 第一条小于 10 分位点的数据
            if lpClose == 0:
                lpClose = row.close
            
            
            lp_decrease = (row.close - lpClose) / lpClose * 100
            lp_decrease = round(lp_decrease, 2)

            item = util_built_item(k, row)
            item['lp_decrease'] = lp_decrease
            item['lp_duration'] = lpDuration
            item['duration'] = len(df[:index])

            ret.loc[len(ret)] = item
    
    return ret


# 月 lp 统计
def lp_static(data):
    ret = None
    for k,df in data.items():
        if ret is None:
            ret = df[['level_point']].copy()

        ret.loc[k] = df['level_point']
    ret = ret.drop(['level_point'], axis=1)

    # 获取最近的月份
    print("排序月份：", ret.T.columns.tolist()[-1])

    tmp = ret.T
    sortKey = tmp.columns.tolist()[-1]
    tmp = tmp.sort_values(by=sortKey)
    
    return tmp.T

# gt1200d 统计
def gt1200d_static(data):
    ret = None
    for k,df in data.items():
        if ret is None:
            ret = df[['gt1200d']].copy()

        ret.loc[k] = df['gt1200d']
    ret = ret.drop(['gt1200d'], axis=1)

    # 按日期倒序
    ret = ret[::-1]

    # 获取最近的月份
    print("排序月份：", ret.T.columns.tolist()[0])
    
    tmp = ret.T
    sortKey = tmp.columns.tolist()[0]
    tmp = tmp.sort_values(by=sortKey, ascending=False)

    return tmp.T


# 月涨跌统计
def pct_static(data):
    ret = None
    for k,df in data.items():
        if ret is None:
            ret = df[['month_pct']]

        ret[k] = df['month_pct']
    
    return ret

# 股息率
def dividend_yield(data):
    df = lib.lib_get_stocks_data()[['trade_code','name','wg','tag','category','bonus_ratio_rmb']]
    df['dividend_yield'] = 0
    df['lp'] = 0

    for index, row in df.iterrows():
        if row['category'] != '股票':
            continue

        stock = row['name']
        bonus_ratio_rmb  = row['bonus_ratio_rmb']
        current_price = data[stock].iloc[-1]['close']
        dividend_yield = round(bonus_ratio_rmb / current_price * 10, 2)
        
        df.loc[index,'dividend_yield'] = dividend_yield
        df.loc[index,'lp'] = data[stock].iloc[-1]['level_point']
    
    return df
